import type { OverWrite } from '@/types';
import type { TableProps } from 'antd';
import { Table } from 'antd';
import React, { useCallback, useContext, useEffect, useMemo, useRef, useState } from 'react';
import { v4 as uuid } from 'uuid';

type Props<RecordType> = OverWrite<
  TableProps<RecordType>,
  {
    scroll: {
      x?: number;
      y: number;
    };
    itemSize: number;
    overscanCount?: number;
  }
>;

interface VirtualTableContextStore {
  clientHeight: number;
  totalHeight: number;
  itemSize: number;
  overscanCount: number;
  dataSourceLength: number;
}

interface ScrollTableContextStore {
  startIndex: number;
  endIndex: number;
}

const ProxyResource: Record<
  string,
  {
    virtualTableContext: React.Context<VirtualTableContextStore>;
    scrollTableContext: React.Context<ScrollTableContextStore>;
  }
> = {};

const CollectionMap = new Proxy(ProxyResource, {
  get(target, virtualContextId: string) {
    if (!target[virtualContextId]) {
      // eslint-disable-next-line no-param-reassign
      target[virtualContextId] = {
        virtualTableContext: React.createContext({
          clientHeight: 0,
          totalHeight: 0,
          itemSize: 0,
          overscanCount: 0,
          dataSourceLength: 0,
        }),
        scrollTableContext: React.createContext({
          startIndex: 0,
          endIndex: 0,
        }),
      };
    }
    return target[virtualContextId];
  },
});

export default function VirtualTable<RecordType extends object = any>(props: Props<RecordType>) {
  const { itemSize, overscanCount = 2, ...restProps } = props;
  const { scroll, dataSource = [] } = restProps;
  const clientHeight = scroll.y;
  const totalHeight = dataSource.length * itemSize;
  const [virtualContextId] = useState(uuid);
  const VirtualTableContext = useMemo(
    () => CollectionMap[virtualContextId].virtualTableContext,
    [virtualContextId],
  );

  const tableComponent: typeof TableComponent = useCallback(
    (props) => <TableComponent {...props} virtualContextId={virtualContextId} />,
    [virtualContextId],
  );

  const bodyWrapperComponent: typeof BodyWrapperComponent = useCallback(
    (props) => <BodyWrapperComponent {...props} virtualContextId={virtualContextId} />,
    [virtualContextId],
  );

  useEffect(
    () => () => {
      delete ProxyResource[virtualContextId];
    },
    [virtualContextId],
  );

  return (
    <VirtualTableContext.Provider
      value={{
        clientHeight,
        totalHeight,
        itemSize,
        overscanCount,
        dataSourceLength: dataSource.length,
      }}
    >
      <Table
        {...restProps}
        components={{
          table: tableComponent,
          body: {
            wrapper: bodyWrapperComponent,
          },
        }}
      />
    </VirtualTableContext.Provider>
  );
}

export function TableComponent(props: {
  style: React.CSSProperties;
  children: React.ReactElement[];
  virtualContextId: string;
}) {
  const { style, children, virtualContextId } = props;
  const scrollRef = useRef<HTMLDivElement>(null);
  const { virtualTableContext: VirtualTableContext, scrollTableContext: ScrollTableContext } =
    useMemo(() => CollectionMap[virtualContextId], [virtualContextId]);
  const { totalHeight, itemSize, overscanCount, clientHeight } = useContext(VirtualTableContext);
  const scrollCollectionRef = useRef<{
    timer: null | number;
  }>({
    timer: null,
  });

  const [scrollTableValue, setScrollTableValue] = useState({
    start: 0,
    scrollTop: 0,
  });

  const { startIndex, endIndex, translateY } = useMemo(() => {
    const startIndex = Math.max(0, scrollTableValue.start - overscanCount);
    const clientShowLength = Math.ceil(clientHeight / itemSize);
    const translateY = startIndex * itemSize;
    return {
      startIndex,
      endIndex: startIndex + overscanCount + clientShowLength - 1 + overscanCount,
      translateY,
    };
  }, [clientHeight, itemSize, overscanCount, scrollTableValue.start]);

  useEffect(() => {
    // 滚动元素
    const scrollElement = scrollRef?.current?.parentElement;
    function onScroll(event: Event) {
      const currentTarget = event.currentTarget as HTMLDivElement;
      // 取消上一次的触发
      if (scrollCollectionRef.current.timer) {
        window.cancelAnimationFrame(scrollCollectionRef.current.timer);
      }

      // 控制更新频率
      scrollCollectionRef.current.timer = window.requestAnimationFrame(() => {
        const { scrollTop } = currentTarget;
        setScrollTableValue((value) => {
          if (value.scrollTop === scrollTop) return value;
          return { start: Math.floor(scrollTop / itemSize), scrollTop };
        });
      });
    }

    scrollElement?.addEventListener('scroll', onScroll);
    return () => {
      scrollElement?.removeEventListener('scroll', onScroll);
    };
  }, [itemSize]);

  return (
    <ScrollTableContext.Provider value={{ startIndex, endIndex }}>
      <table style={{ tableLayout: 'fixed', ...style, transform: `translateY(${translateY}px)` }}>
        {children}
      </table>
      <div
        style={{ height: totalHeight - (endIndex - startIndex + 1) * itemSize }}
        ref={scrollRef}
      />
    </ScrollTableContext.Provider>
  );
}

function BodyWrapperComponent(props: {
  className: string;
  children: React.ReactElement[];
  virtualContextId: string;
}) {
  const { className, children, virtualContextId } = props;
  const items = Array.isArray(children[1]) ? children[1] : [];
  const ScrollTableContext = useMemo(
    () => CollectionMap[virtualContextId].scrollTableContext,
    [virtualContextId],
  );
  const { startIndex, endIndex } = useContext(ScrollTableContext);
  return (
    <tbody className={className}>
      {children.slice(0, 1)}
      {items.length ? items.slice(startIndex, endIndex + 1) : children.slice(1, 2)}
      {children.slice(2)}
    </tbody>
  );
}
